import copy
import json
import sys
import xml.dom.minidom as minidom
import os
import xml.etree.ElementTree as ET
import json


class AnyLabeling2LabelImgConverter:
    def __init__(self, json_folder, output_dir):
        self.json_folder = json_folder
        self.output_dir = output_dir

    def convert(self):
        json_files = [f for f in os.listdir(self.json_folder) if f.endswith('.json')]

        for json_file in json_files:
            json_path = os.path.join(self.json_folder, json_file)

            with open(json_path, 'r', encoding="utf-8") as f:
                data = json.load(f)

            image_filename = data['imagePath']
            image_width = data['imageWidth']
            image_height = data['imageHeight']

            xml_impl = minidom.getDOMImplementation()
            xml_doc = xml_impl.createDocument(None, 'annotation', None)
            xml_root = xml_doc.documentElement

            folder_elem = xml_doc.createElement('folder')

            # ---------------------------------------------------------------
            # XML结果转换文件 folder属性修改---上级目录
            extracted_value = copy.deepcopy(self.output_dir)
            # 找到最后一个反斜杠的位置
            last_backslash_index = extracted_value.rfind("\\")

            if last_backslash_index != -1:  # 如果找到了反斜杠
                extracted_value = extracted_value[last_backslash_index + 1:]  # 获取反斜杠之后的部分

            folder_elem.appendChild(xml_doc.createTextNode(extracted_value))
            # ---------------------------------------------------------------

            # folder_elem.appendChild(xml_doc.createTextNode(self.output_dir))
            xml_root.appendChild(folder_elem)

            filename_elem = xml_doc.createElement('filename')

            # ---------------------------------------------------------------
            # XML结果转换文件 filename属性修改---图片名称
            image_value = copy.deepcopy(image_filename)
            # 找到最后一个反斜杠的位置
            last_image_backslash_index = image_value.rfind("\\")

            if last_backslash_index != -1:  # 如果找到了反斜杠
                image_value = image_value[last_image_backslash_index + 1:]  # 获取反斜杠之后的部分

            filename_elem.appendChild(xml_doc.createTextNode(image_value))
            # ---------------------------------------------------------------

            # filename_elem.appendChild(xml_doc.createTextNode(image_filename))
            xml_root.appendChild(filename_elem)

            path_elem = xml_doc.createElement('path')
            path_elem.appendChild(xml_doc.createTextNode(os.path.join(self.output_dir, image_filename)))
            xml_root.appendChild(path_elem)

            source_elem = xml_doc.createElement('source')
            database_elem = xml_doc.createElement('database')
            database_elem.appendChild(xml_doc.createTextNode('Unknown'))
            source_elem.appendChild(database_elem)
            xml_root.appendChild(source_elem)

            size_elem = xml_doc.createElement('size')
            width_elem = xml_doc.createElement('width')
            width_elem.appendChild(xml_doc.createTextNode(str(image_width)))
            size_elem.appendChild(width_elem)
            height_elem = xml_doc.createElement('height')
            height_elem.appendChild(xml_doc.createTextNode(str(image_height)))
            size_elem.appendChild(height_elem)
            depth_elem = xml_doc.createElement('depth')
            depth_elem.appendChild(xml_doc.createTextNode('3'))
            size_elem.appendChild(depth_elem)
            xml_root.appendChild(size_elem)

            segmented_elem = xml_doc.createElement('segmented')
            segmented_elem.appendChild(xml_doc.createTextNode('0'))
            xml_root.appendChild(segmented_elem)

            for shape in data['shapes']:
                label = shape['label']
                points = shape['points']
                xmin = min(points[0][0], points[1][0])
                ymin = min(points[0][1], points[1][1])
                xmax = max(points[0][0], points[1][0])
                ymax = max(points[0][1], points[1][1])

                object_elem = xml_doc.createElement('object')
                name_elem = xml_doc.createElement('name')
                name_elem.appendChild(xml_doc.createTextNode(label))
                object_elem.appendChild(name_elem)

                # pose
                pose_elem = xml_doc.createElement("pose")
                pose_elem.appendChild(xml_doc.createTextNode("Unspecified"))
                object_elem.appendChild(pose_elem)
                # 截断情况
                truncated_elem = xml_doc.createElement("truncated")
                truncated_elem.appendChild(xml_doc.createTextNode(str(0)))  # 默认为0，表示未截断
                object_elem.appendChild(truncated_elem)
                # 样本困难度
                difficult_elem = xml_doc.createElement("difficult")
                difficult_elem.appendChild(xml_doc.createTextNode(str(0)))  # 默认为0，表示非困难样本
                object_elem.appendChild(difficult_elem)

                bndbox_elem = xml_doc.createElement('bndbox')
                xmin_elem = xml_doc.createElement('xmin')
                xmin_elem.appendChild(xml_doc.createTextNode(str(int(xmin))))
                bndbox_elem.appendChild(xmin_elem)
                ymin_elem = xml_doc.createElement('ymin')
                ymin_elem.appendChild(xml_doc.createTextNode(str(int(ymin))))
                bndbox_elem.appendChild(ymin_elem)
                xmax_elem = xml_doc.createElement('xmax')
                xmax_elem.appendChild(xml_doc.createTextNode(str(int(xmax))))
                bndbox_elem.appendChild(xmax_elem)
                ymax_elem = xml_doc.createElement('ymax')
                ymax_elem.appendChild(xml_doc.createTextNode(str(int(ymax))))
                bndbox_elem.appendChild(ymax_elem)

                object_elem.appendChild(bndbox_elem)
                xml_root.appendChild(object_elem)

            xml_str = xml_doc.toprettyxml(indent="  ")

            output_path = os.path.join(self.output_dir, os.path.splitext(image_filename)[0] + '.xml')

            with open(output_path, 'w', encoding="utf-8") as f:
                f.write(xml_str)

        print('转换完成!')


# Example usage

arguments = sys.argv
# 检查参数数量
if len(arguments) != 2:
    print("请按照格式输入命令，例如：python script.py E:\\folder")
    sys.exit(1)

# 打印传递给脚本的参数
print("当前转换目录：", arguments[1])
converter = AnyLabeling2LabelImgConverter(arguments[1], arguments[1])
# dev
# converter = AnyLabeling2LabelImgConverter('E:\\安全帽\\嵌套二','E:\\安全帽\\嵌套二')
converter.convert()
